# encoding: utf-8
import traceback
import gc

from .base import create_model_with_name
from .rex.predict import predict_contract_elements
from ..config.error import ControlAssemblyError
from ..config.base import PRINT_TRACEBACK_OPTION


class ModelControl(object):
    def __init__(self, default_model_name):
        self.default_model_name = default_model_name
        self.model = load_model(default_model_name)
        self.is_working = False  # 是否在运行

    def predict_elements(self, infos, default_schema):
        self.is_working = True
        try:
            predict_elements = predict_contract_elements(infos, self.model, default_schema,
                                                         is_relative_index=True, is_switch_data=True)
            result = {
                'success': True,
                'predict_elements': predict_elements,
                'error_msg': None
            }
        except:
            if PRINT_TRACEBACK_OPTION:
                traceback.print_exc()
            error_msg = traceback.format_exc()
            result = {
                'success': False,
                'predict_elements': None,
                'error_msg': error_msg
            }
        finally:
            self.is_working = False
        return result

    def release_model(self):
        """
        释放内存
        :return:
        """
        if self.model is not None:
            del self.model
            gc.collect()
        self.model = None


class ParallelModel(object):
    def __init__(self, model_count, default_model_name):
        pass


def load_model(model_name):
    """
    加载模型
    :return:
    """
    try:
        return create_model_with_name(model_name)
    except:
        raise ControlAssemblyError('模型加载出错')